import cv2
from detector import ObjectDetector
from pyramid import pyramid

def process_image(image_path):
    image = cv2.imread(image_path)
    padding = 50
    # 1. 原图加padding
    padded_image = cv2.copyMakeBorder(image, padding, padding, padding, padding, cv2.BORDER_CONSTANT, value=[0,0,0])
    detector = ObjectDetector()
    scale = 1.5
    min_size = (512, 512)
    all_detections = []

    for i, resized in enumerate(pyramid(padded_image, scale=scale, min_size=min_size)):
        # 2. 检测
        detections = detector.detect(resized)
        # 3. 计算当前层的缩放比例
        ratio_x = padded_image.shape[1] / float(resized.shape[1])
        ratio_y = padded_image.shape[0] / float(resized.shape[0])
        for x1, y1, x2, y2, conf, cls in detections:
            # 4. 先映射回padded_image坐标，再减去padding得到原图坐标
            x1o = int(x1 * ratio_x - padding)
            y1o = int(y1 * ratio_y - padding)
            x2o = int(x2 * ratio_x - padding)
            y2o = int(y2 * ratio_y - padding)
            # 可选：确保坐标不越界
            x1o = max(x1o, 0)
            y1o = max(y1o, 0)
            x2o = min(x2o, image.shape[1] - 1)
            y2o = min(y2o, image.shape[0] - 1)
            all_detections.append((x1o, y1o, x2o, y2o, conf, cls))
            # 也可以在每一层保存检测结果
            cv2.rectangle(resized, (x1, y1), (x2, y2), (0,255,0), 2)
            cv2.putText(resized, str(cls), (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
        cv2.imwrite(f'pyramid_result_{i}.jpg', resized)
        print(f'金字塔第{i}层检测到{len(detections)}个目标')

    # 5. 在原图上画所有框
    for x1, y1, x2, y2, conf, cls in all_detections:
        cv2.rectangle(image, (x1, y1), (x2, y2), (0,255,0), 2)
        cv2.putText(image, str(cls), (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)
    cv2.imwrite('all_detections_on_original.jpg', image)
    print('所有检测框已画在原图并保存为 all_detections_on_original.jpg')

if __name__ == "__main__":
    process_image('image1.jpg')

